use crate::action::record::{
    add_type, get_all_types, insert_record, update_geoip_time, update_stat_time, InsertError,
    RecordEntryGeoIpExt,
};
use crate::fairing::device::CertificateItem;
use crate::fairing::record::RecordDb;
use chrono::Utc;
use futures::{future, StreamExt};
use maxminddb::Reader;
use network_scanner_server_ipc::RecordService;
use network_scanner_server_model::{RecordEntry, RecordType, RecordTypeEntry, RecordTypeId};
use rocket_db_pools::sqlx::{self, Acquire};
use rustls::server::WebPkiClientVerifier;
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr},
    sync::Arc,
    time::Duration,
};
use tarpc::context;
use tarpc::{
    server::{self, incoming::Incoming, Channel},
    tokio_serde::formats::Bincode,
};
use tokio::{
    sync::{mpsc, oneshot},
    time,
};
use tokio_rustls::TlsAcceptor;

#[derive(Clone)]
pub struct RecordServiceServer {
    pub tx: mpsc::Sender<RecordEntry>,
    pub db: &'static RecordDb,
}

#[tarpc::server]
impl RecordService for RecordServiceServer {
    async fn add_entry(self, _: context::Context, entry: RecordEntry) {
        if let Err(err) = self.tx.send(entry).await {
            eprintln!("Failed to insert: {:?}", err);
        }
    }
    async fn get_types(self, _: context::Context) -> Vec<RecordType> {
        get_all_types().await
    }
    async fn add_type(
        self,
        _: context::Context,
        r#type: RecordTypeEntry,
    ) -> Result<RecordTypeId, String> {
        match add_type(self.db, r#type).await {
            Ok(id) => Ok(id),
            Err(err) => {
                eprintln!("Failed to insert type: {:?}", err);
                Err(err.to_string())
            }
        }
    }
}

static BUFFER_SIZE: usize = 10000;
static BUFFER_FLUSH_TIME: Duration = Duration::from_secs(5);

pub async fn insert_task(
    db: &RecordDb,
    geoip_reader: Option<Arc<Reader<Vec<u8>>>>,
    mut rx: mpsc::Receiver<RecordEntry>,
    mut stop_rx: oneshot::Receiver<()>,
) {
    let mut buffer = Vec::with_capacity(BUFFER_SIZE);
    let mut geoip_buffer = geoip_reader.as_ref().map(|_| HashMap::with_capacity(1024));
    let mut timer = time::interval(BUFFER_FLUSH_TIME);

    let mut conn = db.acquire().await.expect("Open database connection failed");

    let mut stop = false;
    while !stop {
        let mut timer_terminated = false;

        tokio::select! {
            Some(data) = rx.recv() => {
                if let Some(geoip_reader) = &geoip_reader {
                    let geoip_buffer = geoip_buffer.as_mut().unwrap();
                    if let Some(geoip_items) = data.get_geoip_item(geoip_reader).await {
                        for item in geoip_items {
                            match geoip_buffer.get_mut(&item) {
                                Some(count) => {
                                    *count += 1;
                                },
                                None => {
                                    geoip_buffer.insert(item, 1);
                                }
                            };
                        }
                    }
                }
                buffer.push(data);
            },
            _ = timer.tick() => {
                timer_terminated = true;
            },
            _ = &mut stop_rx => {
                stop = true;
            }
        };

        let size = buffer.len();
        if size >= BUFFER_SIZE || (size > 0 && timer_terminated) || (size > 0 && stop) {
            if let Err(err) = async {
                let mut tx = conn.begin().await?;

                let now = Utc::now();
                update_stat_time(&mut tx, now, buffer.len().try_into().unwrap()).await?;
                if let Some(geoip_buffer) = &mut geoip_buffer {
                    update_geoip_time(&mut tx, now, geoip_buffer).await?;
                    for count in geoip_buffer.values_mut() {
                        *count = 0;
                    }
                }

                for item in buffer.drain(..) {
                    match insert_record(&mut tx, item).await {
                        Ok(_) => {}
                        Err(err) => match err {
                            InsertError::TypeNotFound(id) => {
                                eprintln!("Type id {} not found", id.0)
                            }
                            InsertError::Database(err) => return Err(err),
                        },
                    }
                }
                tx.commit().await?;
                timer.reset();
                println!("Inserted {} columns", size);
                Ok::<(), sqlx::Error>(())
            }
            .await
            {
                eprintln!("Error happened when inserting: {:?}", err);
            }
        }
    }
}

pub async fn ipc_server_task(
    record_db: &'static RecordDb,
    tx: mpsc::Sender<RecordEntry>,
    rpc_cert: Arc<CertificateItem>,
    ca_cert: Arc<CertificateItem>,
) {
    let server_addr = (IpAddr::V4(Ipv4Addr::UNSPECIFIED), 8088);

    let cert_chain: Vec<CertificateDer> = vec![
        rpc_cert.certificate().der().clone(),
        ca_cert.certificate().der().clone(),
    ];
    let private_key: PrivateKeyDer = rpc_cert
        .key_pair()
        .serialize_der()
        .try_into()
        .expect("Invalid key pair");

    let mut root_cert_store = rustls::RootCertStore::empty();
    root_cert_store
        .add(ca_cert.certificate().der().clone())
        .unwrap();
    let verifier = WebPkiClientVerifier::builder(Arc::new(root_cert_store))
        .allow_unknown_revocation_status()
        .build()
        .expect("Unable to build client verifier");

    let config = rustls::ServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(cert_chain, private_key)
        .expect("Invalid server certificate");
    let acceptor = TlsAcceptor::from(Arc::new(config));

    let mut listener = tarpc_rustls::tls::listen(&server_addr, acceptor, Bincode::default)
        .await
        .unwrap();

    listener.config_mut().max_frame_length(1024usize);
    listener
        .filter_map(|r| {
            future::ready(match r {
                Ok(r) => Some(r),
                Err(err) => {
                    eprintln!("Incoming connection failed: {:?}", err);
                    None
                }
            })
        })
        .map(server::BaseChannel::with_defaults)
        .max_channels_per_key(1, |t| t.transport().peer_addr().unwrap().ip())
        .map(|channel| {
            let server = RecordServiceServer {
                tx: tx.clone(),
                db: record_db,
            };
            channel.execute(server.serve())
        })
        .buffer_unordered(256)
        .for_each(|_| async {})
        .await;
}
